import torch
from torch import nn
from torch.utils.data import DataLoader

from titanic.TitanicDataset import TitanicDataset


# 建立神经网络模型
class NeuralNetwork(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        # 定义多层网络模型
        self.fc1 = nn.Linear(6, 32)
        self.fc2 = nn.Linear(32, 16)
        self.fc3 = nn.Linear(16, 8)
        self.fc4 = nn.Linear(8, 1)
        # 定义激活函数
        self.activation = nn.Sigmoid()

    def forward(self, x):
        x = self.activation(self.fc1(x))
        x = self.activation(self.fc2(x))
        x = self.activation(self.fc3(x))
        x = self.activation(self.fc4(x))
        return x


if __name__ == '__main__':
    # 初始化模型
    model = NeuralNetwork()
    # 定义数据集
    dataset = TitanicDataset("data/train.csv")
    # 定义数据加载
    dataLoader = DataLoader(dataset, batch_size=16, shuffle=True, num_workers=2)

    # 定义损失
    criterion = nn.BCELoss(reduction='mean')
    # 定义优化器
    optimizer = torch.optim.Rprop(model.parameters(), lr=0.01)

    # 记录最低的loss下的权重
    minLoss = 9999
    # 训练
    for epoch in range(200):
        epochLoss = 9999
        for idx, batch in enumerate(dataLoader):
            x, y = batch
            # 计算前馈
            y_pred = model(x)
            # 计算loss
            loss = criterion(y_pred, y)
            # 清除梯度
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            # 记录最好的数据
            if loss.item() < minLoss:
                minLoss = loss.item()
                # 保存当前模型
                torch.save(model.state_dict(), f'models/best.pth')
            epochLoss = loss.item()
        # 打印
        print(f"Epoch: {epoch}, Loss: {epochLoss}, best loss: {minLoss}")
